import btk
import unittest
import _TDDConfigure

class IMUsExtractorTest(unittest.TestCase):
  def test_NoAcquisition(self):
    imuse = btk.btkIMUsExtractor()
    output = imuse.GetOutput()
    output.Update()
    self.assertEqual(output.GetItemNumber(), 0)

  def test_NoIMU(self):
    imuse = btk.btkIMUsExtractor()
    imuse.SetInput(btk.btkAcquisition())
    output = imuse.GetOutput()
    output.Update()
    self.assertEqual(output.GetItemNumber(), 0)

  def test_OneAcquisition(self):
    acq = btk.btkAcquisition()
    acq.Init(0,200,6)
    imu = btk.btkMetaDataCreateChild(acq.GetMetaData(), 'IMU')
    btk.btkMetaDataCreateChild(imu, 'USED', 1)
    btk.btkMetaDataCreateChild(imu, 'LABELS', btk.btkStringArray(1,'Foo'))
    btk.btkMetaDataCreateChild(imu, 'DESCRIPTIONS', btk.btkStringArray(1,'BAR'))
    btk.btkMetaDataCreateChild(imu, 'TYPE', btk.btkIntArray(1,1))
    channels = btk.btkIntArray(6); channels[0] = 1; channels[1] = 2; channels[2] = 3; channels[3] = 4; channels[4] = 5; channels[5] = 6
    btk.btkMetaDataCreateChild(imu, 'CHANNEL', channels, 6)
    imuse = btk.btkIMUsExtractor()
    imuse.SetInput(acq)
    output = imuse.GetOutput()
    output.Update()
    self.assertEqual(output.GetItemNumber(), 1)
    self.assertEqual(output.GetItem(0).GetLabel(), 'Foo')
    self.assertEqual(output.GetItem(0).GetDescription(), 'BAR')
    self.assertEqual(output.GetItem(0).GetAccelerometerX().GetLabel(), acq.GetAnalog(0).GetLabel())
    self.assertEqual(output.GetItem(0).GetAccelerometerY().GetLabel(), acq.GetAnalog(1).GetLabel())
    self.assertEqual(output.GetItem(0).GetAccelerometerZ().GetLabel(), acq.GetAnalog(2).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeX().GetLabel(), acq.GetAnalog(3).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeY().GetLabel(), acq.GetAnalog(4).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeZ().GetLabel(), acq.GetAnalog(5).GetLabel())

  def test_OneAcquisition_NoDescriptions(self):
    acq = btk.btkAcquisition()
    acq.Init(0,200,6)
    imu = btk.btkMetaDataCreateChild(acq.GetMetaData(), 'IMU')
    btk.btkMetaDataCreateChild(imu, 'USED', 1)
    btk.btkMetaDataCreateChild(imu, 'LABELS', btk.btkStringArray(1,'Foo'))
    btk.btkMetaDataCreateChild(imu, 'TYPE', btk.btkIntArray(1,1))
    channels = btk.btkIntArray(6); channels[0] = 6; channels[1] = 5; channels[2] = 4; channels[3] = 3; channels[4] = 2; channels[5] = 1
    btk.btkMetaDataCreateChild(imu, 'CHANNEL', channels, 6)
    imuse = btk.btkIMUsExtractor()
    imuse.SetInput(acq)
    output = imuse.GetOutput()
    output.Update()
    self.assertEqual(output.GetItemNumber(), 1)
    self.assertEqual(output.GetItem(0).GetLabel(), 'Foo')
    self.assertEqual(output.GetItem(0).GetDescription(), '')
    self.assertEqual(output.GetItem(0).GetAccelerometerX().GetLabel(), acq.GetAnalog(5).GetLabel())
    self.assertEqual(output.GetItem(0).GetAccelerometerY().GetLabel(), acq.GetAnalog(4).GetLabel())
    self.assertEqual(output.GetItem(0).GetAccelerometerZ().GetLabel(), acq.GetAnalog(3).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeX().GetLabel(), acq.GetAnalog(2).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeY().GetLabel(), acq.GetAnalog(1).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeZ().GetLabel(), acq.GetAnalog(0).GetLabel())

  def test_OneAcquisition_NoLabelsAndDescriptions(self):
    acq = btk.btkAcquisition()
    acq.Init(0,200,6)
    imu = btk.btkMetaDataCreateChild(acq.GetMetaData(), 'IMU')
    btk.btkMetaDataCreateChild(imu, 'USED', 1)
    btk.btkMetaDataCreateChild(imu, 'TYPE', btk.btkIntArray(1,1))
    channels = btk.btkIntArray(6); channels[0] = 6; channels[1] = 5; channels[2] = 4; channels[3] = 3; channels[4] = 2; channels[5] = 1
    btk.btkMetaDataCreateChild(imu, 'CHANNEL', channels, 6)
    imuse = btk.btkIMUsExtractor()
    imuse.SetInput(acq)
    output = imuse.GetOutput()
    output.Update()
    self.assertEqual(output.GetItemNumber(), 1)
    self.assertEqual(output.GetItem(0).GetLabel(), 'IMU #1')
    self.assertEqual(output.GetItem(0).GetDescription(), '')
    self.assertEqual(output.GetItem(0).GetAccelerometerX().GetLabel(), acq.GetAnalog(5).GetLabel())
    self.assertEqual(output.GetItem(0).GetAccelerometerY().GetLabel(), acq.GetAnalog(4).GetLabel())
    self.assertEqual(output.GetItem(0).GetAccelerometerZ().GetLabel(), acq.GetAnalog(3).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeX().GetLabel(), acq.GetAnalog(2).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeY().GetLabel(), acq.GetAnalog(1).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeZ().GetLabel(), acq.GetAnalog(0).GetLabel())

  def test_OneAcquisition_MissingChannels(self):
    acq = btk.btkAcquisition()
    acq.Init(0,200,6)
    imu = btk.btkMetaDataCreateChild(acq.GetMetaData(), 'IMU')
    btk.btkMetaDataCreateChild(imu, 'USED', 1)
    btk.btkMetaDataCreateChild(imu, 'TYPE', btk.btkIntArray(1,1))
    channels = btk.btkIntArray(6); channels[0] = 6; channels[1] = 5; channels[2] = -1; channels[3] = -1; channels[4] = 2; channels[5] = 1
    btk.btkMetaDataCreateChild(imu, 'CHANNEL', channels, 6)
    imuse = btk.btkIMUsExtractor()
    imuse.SetInput(acq)
    output = imuse.GetOutput()
    output.Update()
    self.assertEqual(output.GetItemNumber(), 1)
    self.assertEqual(output.GetItem(0).GetLabel(), 'IMU #1')
    self.assertEqual(output.GetItem(0).GetDescription(), '')
    self.assertEqual(output.GetItem(0).GetAccelerometerX().GetLabel(), acq.GetAnalog(5).GetLabel())
    self.assertEqual(output.GetItem(0).GetAccelerometerY().GetLabel(), acq.GetAnalog(4).GetLabel())
    # self.assertTrue(output.GetItem(0).GetAccelerometerZ()._get() == 0)
    # self.assertTrue(output.GetItem(0).GetGyroscopeX()._get() == 0)
    self.assertEqual(output.GetItem(0).GetGyroscopeY().GetLabel(), acq.GetAnalog(1).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeZ().GetLabel(), acq.GetAnalog(0).GetLabel())

  def test_OneAcquisition_TwoSensors(self):
    acq = btk.btkAcquisition()
    acq.Init(0,200,15)
    imu = btk.btkMetaDataCreateChild(acq.GetMetaData(), 'IMU')
    btk.btkMetaDataCreateChild(imu, 'USED', 2)
    btk.btkMetaDataCreateChild(imu, 'LABELS', btk.btkStringArray(2,'Foo'))
    btk.btkMetaDataCreateChild(imu, 'TYPE', btk.btkIntArray(2,1))
    channels = btk.btkIntArray(18);
    channels[0] = 1; channels[1] = 2; channels[2] = 3; channels[3] = 4; channels[4] = 5; channels[5] = 6; channels[6] = -1; channels[7] = -1; channels[8] = -1
    channels[9] = 7; channels[10] = 8; channels[11] = 9; channels[12] = 10; channels[13] = 11; channels[14] = 12; channels[15] = 13; channels[16] = 14; channels[17] = 15
    btk.btkMetaDataCreateChild(imu, 'CHANNEL', channels, 9)
    extra = btk.btkIntArray(2); extra[0] = 0; extra[1] = 3;
    btk.btkMetaDataCreateChild(imu, 'EXTRA', extra)
    imuse = btk.btkIMUsExtractor()
    imuse.SetInput(acq)
    output = imuse.GetOutput()
    output.Update()
    self.assertEqual(output.GetItemNumber(), 2)
    self.assertEqual(output.GetItem(0).GetLabel(), 'Foo')
    self.assertEqual(output.GetItem(0).GetDescription(), '')
    self.assertEqual(output.GetItem(0).GetAccelerometerX().GetLabel(), acq.GetAnalog(0).GetLabel())
    self.assertEqual(output.GetItem(0).GetAccelerometerY().GetLabel(), acq.GetAnalog(1).GetLabel())
    self.assertEqual(output.GetItem(0).GetAccelerometerZ().GetLabel(), acq.GetAnalog(2).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeX().GetLabel(), acq.GetAnalog(3).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeY().GetLabel(), acq.GetAnalog(4).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeZ().GetLabel(), acq.GetAnalog(5).GetLabel())
    self.assertEqual(output.GetItem(1).GetLabel(), 'Foo')
    self.assertEqual(output.GetItem(1).GetDescription(), '')
    self.assertEqual(output.GetItem(1).GetAccelerometerX().GetLabel(), acq.GetAnalog(6).GetLabel())
    self.assertEqual(output.GetItem(1).GetAccelerometerY().GetLabel(), acq.GetAnalog(7).GetLabel())
    self.assertEqual(output.GetItem(1).GetAccelerometerZ().GetLabel(), acq.GetAnalog(8).GetLabel())
    self.assertEqual(output.GetItem(1).GetGyroscopeX().GetLabel(), acq.GetAnalog(9).GetLabel())
    self.assertEqual(output.GetItem(1).GetGyroscopeY().GetLabel(), acq.GetAnalog(10).GetLabel())
    self.assertEqual(output.GetItem(1).GetGyroscopeZ().GetLabel(), acq.GetAnalog(11).GetLabel())
    self.assertEqual(output.GetItem(1).GetChannel(6).GetLabel(), acq.GetAnalog(12).GetLabel())
    self.assertEqual(output.GetItem(1).GetChannel(7).GetLabel(), acq.GetAnalog(13).GetLabel())
    self.assertEqual(output.GetItem(1).GetChannel(8).GetLabel(), acq.GetAnalog(14).GetLabel())

  def test_TwoAcquisitions(self):
    # Acquisition #1
    acq1 = btk.btkAcquisition()
    acq1.Init(0,200,6)
    imu = btk.btkMetaDataCreateChild(acq1.GetMetaData(), 'IMU')
    btk.btkMetaDataCreateChild(imu, 'USED', 1)
    btk.btkMetaDataCreateChild(imu, 'DESCRIPTIONS', btk.btkStringArray(1,'BAR'))
    btk.btkMetaDataCreateChild(imu, 'TYPE', btk.btkIntArray(1,1))
    channels = btk.btkIntArray(6); channels[0] = 1; channels[1] = 2; channels[2] = 3; channels[3] = 4; channels[4] = 5; channels[5] = 6
    btk.btkMetaDataCreateChild(imu, 'CHANNEL', channels, 6)
    # Acquisition #2 (two sensors)
    acq2 = btk.btkAcquisition()
    acq2.Init(0,200,15)
    imu = btk.btkMetaDataCreateChild(acq2.GetMetaData(), 'IMU')
    btk.btkMetaDataCreateChild(imu, 'USED', 2)
    btk.btkMetaDataCreateChild(imu, 'LABELS', btk.btkStringArray(2,'Foo'))
    btk.btkMetaDataCreateChild(imu, 'TYPE', btk.btkIntArray(2,1))
    channels.resize(18) 
    channels[0] = 1; channels[1] = 2; channels[2] = 3; channels[3] = 4; channels[4] = 5; channels[5] = 6; channels[6] = -1; channels[7] = -1; channels[8] = -1
    channels[9] = 7; channels[10] = 8; channels[11] = 9; channels[12] = 10; channels[13] = 11; channels[14] = 12; channels[15] = 13; channels[16] = 14; channels[17] = 15
    btk.btkMetaDataCreateChild(imu, 'CHANNEL', channels, 9)
    extra = btk.btkIntArray(2); extra[0] = 0; extra[1] = 3;
    btk.btkMetaDataCreateChild(imu, 'EXTRA', extra)
    imuse = btk.btkIMUsExtractor()
    imuse.SetInput(0, acq1)
    imuse.SetInput(1, acq2)
    output = imuse.GetOutput()
    output.Update()
    self.assertEqual(output.GetItemNumber(), 3)
    self.assertEqual(output.GetItem(0).GetLabel(), 'IMU #1')
    self.assertEqual(output.GetItem(0).GetDescription(), 'BAR')
    self.assertEqual(output.GetItem(0).GetAccelerometerX().GetLabel(), acq1.GetAnalog(0).GetLabel())
    self.assertEqual(output.GetItem(0).GetAccelerometerY().GetLabel(), acq1.GetAnalog(1).GetLabel())
    self.assertEqual(output.GetItem(0).GetAccelerometerZ().GetLabel(), acq1.GetAnalog(2).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeX().GetLabel(), acq1.GetAnalog(3).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeY().GetLabel(), acq1.GetAnalog(4).GetLabel())
    self.assertEqual(output.GetItem(0).GetGyroscopeZ().GetLabel(), acq1.GetAnalog(5).GetLabel())
    self.assertEqual(output.GetItem(1).GetLabel(), 'Foo')
    self.assertEqual(output.GetItem(1).GetDescription(), '')
    self.assertEqual(output.GetItem(1).GetAccelerometerX().GetLabel(), acq2.GetAnalog(0).GetLabel())
    self.assertEqual(output.GetItem(1).GetAccelerometerY().GetLabel(), acq2.GetAnalog(1).GetLabel())
    self.assertEqual(output.GetItem(1).GetAccelerometerZ().GetLabel(), acq2.GetAnalog(2).GetLabel())
    self.assertEqual(output.GetItem(1).GetGyroscopeX().GetLabel(), acq2.GetAnalog(3).GetLabel())
    self.assertEqual(output.GetItem(1).GetGyroscopeY().GetLabel(), acq2.GetAnalog(4).GetLabel())
    self.assertEqual(output.GetItem(1).GetGyroscopeZ().GetLabel(), acq2.GetAnalog(5).GetLabel())
    self.assertEqual(output.GetItem(2).GetLabel(), 'Foo')
    self.assertEqual(output.GetItem(2).GetDescription(), '')
    self.assertEqual(output.GetItem(2).GetAccelerometerX().GetLabel(), acq2.GetAnalog(6).GetLabel())
    self.assertEqual(output.GetItem(2).GetAccelerometerY().GetLabel(), acq2.GetAnalog(7).GetLabel())
    self.assertEqual(output.GetItem(2).GetAccelerometerZ().GetLabel(), acq2.GetAnalog(8).GetLabel())
    self.assertEqual(output.GetItem(2).GetGyroscopeX().GetLabel(), acq2.GetAnalog(9).GetLabel())
    self.assertEqual(output.GetItem(2).GetGyroscopeY().GetLabel(), acq2.GetAnalog(10).GetLabel())
    self.assertEqual(output.GetItem(2).GetGyroscopeZ().GetLabel(), acq2.GetAnalog(11).GetLabel())
    self.assertEqual(output.GetItem(2).GetChannel(6).GetLabel(), acq2.GetAnalog(12).GetLabel())
    self.assertEqual(output.GetItem(2).GetChannel(7).GetLabel(), acq2.GetAnalog(13).GetLabel())
    self.assertEqual(output.GetItem(2).GetChannel(8).GetLabel(), acq2.GetAnalog(14).GetLabel())

  def test_TwoAcquisitions_NoIMU(self):
    imuse = btk.btkIMUsExtractor()
    imuse.SetInput(0,btk.btkAcquisition())
    imuse.SetInput(1,btk.btkAcquisition())
    output = imuse.GetOutput()
    output.Update()
    self.assertEqual(output.GetItemNumber(), 0)